from playwright.async_api import async_playwright,Browser, BrowserContext, Page,ConsoleMessage
from pydantic import BaseModel,Field
import asyncio
import os
from typing import Set, Optional
from urllib.parse import urlparse
import json  
from llmweb.logging import logger

class BrowserPool: 
    """
    浏览器缓存池，负责管理Page的创建与销毁
    # 示例用法
    browser_pool = BrowserPool()
    await browser_pool.start_periodic_context_cleanup()

    # ... 在适当的时候关闭浏览器池和取消任务 ...
    await browser_pool.close()
    if browser_pool.close_stale_contexts_task:
        await browser_pool.close_stale_contexts_task  # 等待任务完成或取消
    """
    browser: Browser
    context: BrowserContext
    close_stale_contexts_task: asyncio.Task = None
    headless: bool = True
    workdir: str = "."
    pool_size: int
    page_pool: asyncio.Queue

    def __init__(self, headless: Optional[bool] = True, 
                 browser: Optional[Browser] = None,
                 workdir: Optional[str] = None,
                 pool_size: int = 10) -> None:
        """
        初始化一个浏览器实例。
        
        参数:
        headless: 可选，是否以无头模式启动浏览器，默认为True。
        browser: 可选，指定使用的浏览器类型，默认为None，会根据环境自动选择。
        workdir: 可选，设置工作目录，默认为None，如果设置，则使用指定目录，否则使用当前工作目录。
        pool_size: int，页面池的大小，默认为10。
        
        返回值:
        无
        """
        self.headless = headless
        self.browser = browser
        self.context = None
        if workdir is not None:
            self.workdir = workdir
        else:
            self.workdir = os.getcwd()
        self.pool_size = pool_size
        self.page_pool = None

    async def initialize(self):  
        if self.page_pool is None:  
            self.page_pool = asyncio.Queue(maxsize=self.pool_size)  
            for _ in range(self.pool_size):  
                await self.page_pool.put(await self.create_page())

    async def create_page(self):
        """
        异步创建一个新页面。
        
        参数:
        self - 对象本身的引用。
        
        返回值:
        page - 新创建的页面对象。
        """
        ctx = await self.get_context()
        page = await ctx.new_page()
        return page

    async def acquire_page(self):
        obj = await self.page_pool.get()  
        return obj  
    
    async def release_page(self, page: Page):
        if self.page_pool.qsize() < self.pool_size:  
            await self.page_pool.put(page)
    
    async def get_browser(self) -> Browser:
        """
        异步获取一个浏览器实例。
        
        如果浏览器实例还未创建，则尝试启动一个Firefox浏览器实例（根据playwright的配置）。
        如果启动失败，记录错误并抛出异常。
        
        返回值:
            返回一个Browser类型的实例，表示已启动的浏览器。
        """
        if not self.browser:
            try:
                playwright = await async_playwright().start()
                self.browser = await playwright.firefox.launch(headless=self.headless)
            except Exception as e:
                logger.error(f"Failed to launch browser: {e}")
                await playwright.stop()
                raise e
        return self.browser

    async def get_context(self) -> BrowserContext:
        
        br = await self.get_browser()
        self.context = await br.new_context()
        return self.context
    
    @staticmethod
    async def get_domain(page: Page) -> str:
        parsed_url = urlparse(page.url)
        return parsed_url.netloc
    
    @staticmethod
    async def get_local_storage(page: Page):
        # 使用 evaluate 方法执行 JavaScript 代码来获取 localStorage 的内容
        storage = await page.evaluate("() => Object.assign({}, window.localStorage)")
        return storage
    
    async def save_cookies(self, page: Page) -> None:
        domain = await self.get_domain(page)
        cookies_path = os.path.join(self.workdir, f"{domain}_cookies.json")
        cookies = await page.context.cookies()
        with open(cookies_path, "w", encoding="utf-8") as f:
            ck = json.dumps(cookies)
            f.write(ck)
        
    async def load_cookies(self, page: Page) -> None:
        domain = await self.get_domain(page)
        cookies_path = os.path.join(self.workdir, f"{domain}_cookies.json") 
        try:
            with open(cookies_path, "r", encoding="utf-8") as f:
                cookies = json.load(f)
                await page.context.clear_cookies()
                await page.context.add_cookies(cookies)
        except Exception as e:
            logger.error(f"Failed to load cookies: {e}")

    async def save_local_storage(self, page: Page) -> None:
        domain = await self.get_domain(page)
        local_storage_path = os.path.join(self.workdir, f"{domain}_local_storage.json")
        storage = await page.evaluate("() => Object.assign({}, window.localStorage)")
        with open(local_storage_path, "w", encoding="utf-8") as f:
            ck = json.dumps(storage)
            f.write(ck)
    
    async def load_local_storage(self, page: Page) -> None:
        domain = await self.get_domain(page)
        local_storage_path = os.path.join(self.workdir, f"{domain}_local_storage.json")
        try:
            with open(local_storage_path, 'r', encoding='utf-8') as f:
                storage_json = f.read()
                # 将JSON字符串转换回Python字典
                storage = json.loads(storage_json)
                for key, value in storage.items():
                    logger.debug(f"locad localStorage key: {key}, value: {value}")
                    await page.evaluate(f"([key, value]) => window.localStorage.setItem(key, value)", [key, value])
        except Exception as e:
            logger.error(f"Failed to load localStorage: {e}")
    
    async def close_stale_contexts(self, interval_minutes: int = 10) -> None:
        while True:
            try:
                await asyncio.sleep(interval_minutes * 60)
                if not self.context.is_traced_recently():
                    await self.context.close()
                    await self.browser.close()
            except asyncio.CancelledError:
                break  # 任务被取消，退出循环
    
    async def close(self):
        if self.context:
            await self.context.close()
            self.context = None
        if self.browser:
            await self.browser.close()
            self.browser = None
        if self.close_stale_contexts_task and not self.close_stale_contexts_task.done():
            self.close_stale_contexts_task.cancel()  # 取消任务
        while not self.page_pool.empty():  
            await self.page_pool.get()

    async def start_periodic_context_cleanup(self):
        self.close_stale_contexts_task = asyncio.create_task(self.close_stale_contexts())
        return self.close_stale_contexts_task


    async def add_save_cookie_btn(self, page: Page):
        await page.wait_for_load_state('networkidle')
        await page.eval_on_selector('body', """  
            var button = document.createElement('button');
            button.id = 'llmweb_save_cookies';
            button.textContent = 'Save Cookies';
            // 设置button的样式  
            button.style.position = "fixed"; // 固定位置  
            button.style.top = "10px"; // 距离页面顶部的距离  
            button.style.right = "10px"; // 距离页面右侧的距离  
            button.style.zIndex = "9999"; // 确保按钮在其他元素之上  
            button.style.backgroundColor = "red"; // 设置背景颜色  
            button.style.color = "white"; // 设置文本颜色  
            button.style.border = "none"; // 无边框  
            button.style.padding = "10px 20px"; // 设置内边距  
            button.style.cursor = "pointer"; // 鼠标悬停时显示指针    
            document.body.appendChild(button);
            button.addEventListener('click', function() {  
                console.log('SAVE_COOKIES');  
            });  
        """)
        async def on_message(message: ConsoleMessage):
            if 'SAVE_COOKIES' in message.text:  
                await self.save_cookies(page)
                await self.save_local_storage(page)

        page.on('console', on_message)
